Window function fusion #29276

ggevay · 2024-08-29T11:31:14Z

This adds the fusion optimization for value window functions (lag, lead, first_value, last_value), which is badly needed by https://github.com/MaterializeInc/accounts/issues/3

The point is that it amortizes the various overheads of window functions (e.g., the wrapping/unwrapping of other columns, which gets quadratically big with the number of window functions) by performing the MIR window function pattern only once for a group of window functions.

I've added this as an HIR transform, because an MIR transform would have needed to recognize and manipulate the complicated MIR window function pattern.

See the doc comment on fuse_window_functions for how we choose the groups of window functions to be fused.

I've added a new value window function called Fused, and implemented the corresponding MIR AggregateFunc. In several places in the implementation of AggregateFunc::FusedValueWindowFunc, I refactored the implementations of existing value window functions in order to be able to reuse as much code from them as possible.

Note that this does not yet implement fusion for window aggregations, which would be needed by https://github.com/MaterializeInc/accounts/issues/77. I'll implement that in a follow-up PR soon.

Motivation

This PR adds a known-desirable feature: https://github.com/MaterializeInc/database-issues/issues/6247

Tips for reviewer

Look at the commits separately. The first two are just minor preparations (see commit msgs), the 3rd commit is the main one, and then there is a commit that guards the new code behind a feature flag.

Checklist

This PR has adequate test coverage / QA involvement has been duly considered. (trigger-ci for additional test/nightly runs)
- Added some new slt tests.
- In the already-existing tests in window_funcs.slt, there were 20 cases where fusion happened.
- There is also an RQG grammar specifically for window functions.
  - First run: https://buildkite.com/materialize/nightly/builds/9307#01919def-9f3d-45e6-8651-d03339f6d5ce
  - Second run: https://buildkite.com/materialize/nightly/builds/9309
This PR has an associated up-to-date design doc, is a design doc (template), or is sufficiently small to not require a design.
If this PR evolves an existing $T ⇔ Proto$T mapping (possibly in a backwards-incompatible way), then it is tagged with a T-proto label.
If this PR will require changes to cloud orchestration or tests, there is a companion cloud PR to account for those changes that is tagged with the release-blocker label (example).
If this PR includes major user-facing behavior changes, I have pinged the relevant PM to schedule a changelog post.

I kept forgetting these matches when adding new aggregate functions. Now the compiler will tell me if I forget one of these.

For the rename, see https://materializeinc.slack.com/archives/C02PPB50ZHS/p1699536714967029

jkosh44

adapter changes LGTM

petrosagg · 2024-08-29T14:16:05Z

src/expr/src/relation.rs

+                    .clone();
+
+                let all_func_return_types =
+                    return_type_with_orig_row.unwrap_record_element_type()[0].clone();


return_type_with_orig_row.unwrap_record_element_type() which is also used below could be calculated once and stored in a variable

Below it is slightly different: it's all_func_return_types.unwrap_record_element_type()

src/expr/src/relation/func.rs

petrosagg · 2024-08-29T15:34:47Z

src/expr/src/relation/func.rs

+/// The expected input is in the format of `[((OriginalRow, (Args1, Args2, ...)), OrderByExprs...)]`
+/// where `Args1`, `Args2`, are the arguments of each of the fused functions. For functions that
+/// have only a single argument (first_value/last_value), these are simple values. For functions
+/// that have multiple arguments (lag/lead), these are also records.


Is it ok that these values will have a type that is not expressible by materialize today (effecticely a union)?

Yes, it's true that this function can handle multiple different types at each argument, and this couldn't be represented by our type system. However, the union fortunately doesn't need to be represented in the type system, because we know the actual type when looking at a specific query.

One way to think about this situation is that AggregateFunc::FusedValueWindowFunc is kind of a template function, in that it's parameterized by the functions that got fused. fused_value_window_func is an implementation of all possible instantiations of this template function.

petrosagg · 2024-08-29T15:35:52Z

src/expr/src/relation/func.rs

+    // Let's create a new RowArena, to avoid flooding the caller's arena with stuff proportional to
+    // the window partition size. We will use our own arena for internal evaluations, and will use
+    // the caller's at the very end to return the final result Datum.
+    let temp_storage = RowArena::new();


What goes wrong if we use the caller's arena? Won't we deallocate it very soon anyway?

Well, might be ok, but I'm not sure. There are many callers, and I didn't want to check how soon each of them deallocates it.

(window_aggr does the same thing.)

petrosagg · 2024-08-29T15:36:31Z

src/expr/src/relation/func.rs

+
+    let input_datums_with_ranks = order_aggregate_datums_with_rank(input_datums, order_by);
+
+    let mut encoded_argsss = vec![Vec::new(); funcs.len()];


Is this a typo? or does the sss signify something?

This is a Haskell naming convention for keeping track of collection nesting depth: when naming a nested collection, we add an s at each level. Like,

if I have a variable that holds a single argument, it's named arg,

if I have a variable that is a collection of arguments, then it's args,

if I have a variable that is a collection of collections of arguments, then it's argss,

if I have a variable that is a collection of collections of collections of arguments, then it's argsss.

The last one is the case here, because encoded_argsss is a Vec<Vec<Datum>>, where each Datum can be a record of multiple arguments. (The outer Vec is over input rows, the inner Vec is over the different constituents of the fused function, and the Datum is the arguments to a single constituent function, e.g., lag's 3 arguments.)

Haskell example: https://stackoverflow.com/questions/66066788/i-am-struggling-to-understand-how-concat-is-read-in-haskell

petrosagg · 2024-08-29T15:45:11Z

src/sql/src/plan/expr.rs

+                order_by,
+                window_frame,
+            },
+            ValueWindowFunc::LastValue => mz_expr::AggregateFunc::LastValue {


Not related to this PR but why is LagLead folded into a single aggregate but FirstValue and LastValue are different?

This predates my time at Materialize, but indeed, it looks like it might be better to fold this into a FirstLastValue.

petrosagg · 2024-08-29T15:46:50Z

src/sql/src/plan/expr.rs

+    /// references that are at a deeper subquery nesting depth, but refer back to the root level.
+    /// (Note that even if `self` is embedded inside a larger expression, we consider the
+    /// "root level" to be `self`'s level.)
+    pub fn visit_columns_referring_to_root_level<F>(&self, f: &mut F)


Can you give an example of the columns that this visits? What does the usize passed to the visitor function represent?

This visits all column references in the expression, at every level.

The usize given to the visitor is the subquery nesting depth: When recursively descending down the expression, if we are descending into a subquery, we increase the nesting depth.

(Also discussed during the call.)

petrosagg · 2024-08-29T15:48:09Z

src/sql/src/plan/lowering.rs

@@ -455,6 +458,7 @@ impl HirRelationExpr {
                                requires_nonexistent_column
                            })
                            .unwrap_or(scalars.len());
+                        assert!(end_idx > 0, "a Map expression references itself or a later column; lowered_arity: {}, expressions: {:?}", lowered_arity, scalars);


Is this related to this PR or a thing that would be useful independently?

Useful independently. Without the assert, if there is a bug in the preparation of the HIR expressions such that this condition ends up false, then we get into an infinite loop here, which would be really bad in production. It's much less bad to crash instead. (Without this assert, an envd thread would be spinning forever, and then it would be really hard to debug what went wrong if this happens in production. But even when it happens locally it makes it easier to figure out what's going on.)

During the development of this PR, I twice had bugs where I got into an infinite loop here. I can imagine similar bugs happening in other situations later.

petrosagg

Went over this a second time with Gabor and everything LGTM!

src/expr/src/relation/func.rs

ggevay · 2024-08-29T17:46:51Z

Thank you very much for the reviews! Merging!

ggevay added 2 commits August 29, 2024 08:26

Remove wildcard patterns from AggregateFunc metadata functions

1bfcffa

I kept forgetting these matches when adding new aggregate functions. Now the compiler will tell me if I forget one of these.

Minor refactorings around window functions

8deae07

For the rename, see https://materializeinc.slack.com/archives/C02PPB50ZHS/p1699536714967029

ggevay added A-optimization Area: query optimization and transformation A-CLUSTER Topics related to the CLUSTER layer labels Aug 29, 2024

ggevay force-pushed the win-fusion-value branch 3 times, most recently from 96b89c9 to 9611c21 Compare August 29, 2024 13:29

ggevay marked this pull request as ready for review August 29, 2024 13:29

ggevay requested review from a team as code owners August 29, 2024 13:29

ggevay requested review from jkosh44 and petrosagg August 29, 2024 13:29

jkosh44 approved these changes Aug 29, 2024

View reviewed changes

petrosagg reviewed Aug 29, 2024

View reviewed changes

petrosagg approved these changes Aug 29, 2024

View reviewed changes

src/expr/src/relation/func.rs Show resolved Hide resolved

ggevay added 2 commits August 29, 2024 19:45

Fusion for value window functions

cee22a7

Add feature flag for value window function fusion

b1d8778

ggevay force-pushed the win-fusion-value branch from 9611c21 to b1d8778 Compare August 29, 2024 17:45

ggevay enabled auto-merge August 29, 2024 17:46

ggevay merged commit 9d6ade0 into MaterializeInc:main Aug 29, 2024
81 checks passed

github-actions bot locked and limited conversation to collaborators Aug 29, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Window function fusion #29276

Window function fusion #29276

ggevay commented Aug 29, 2024 •

edited

Loading

jkosh44 left a comment

petrosagg Aug 29, 2024

ggevay Aug 29, 2024

petrosagg Aug 29, 2024

ggevay Aug 29, 2024

petrosagg Aug 29, 2024

ggevay Aug 29, 2024

petrosagg Aug 29, 2024

ggevay Aug 29, 2024

petrosagg Aug 29, 2024

ggevay Aug 29, 2024

petrosagg Aug 29, 2024

ggevay Aug 29, 2024

petrosagg Aug 29, 2024

ggevay Aug 29, 2024

petrosagg left a comment

ggevay commented Aug 29, 2024


		let input_datums_with_ranks = order_aggregate_datums_with_rank(input_datums, order_by);

		let mut encoded_argsss = vec![Vec::new(); funcs.len()];

Window function fusion #29276

Window function fusion #29276

Conversation

ggevay commented Aug 29, 2024 • edited Loading

Motivation

Tips for reviewer

Checklist

jkosh44 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

petrosagg left a comment

Choose a reason for hiding this comment

ggevay commented Aug 29, 2024

ggevay commented Aug 29, 2024 •

edited

Loading